#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QTableWidget, QTableWidgetItem,
                           QHeaderView, QAbstractItemView, QMenu, QAction)
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtGui import QColor, QBrush

class PacketListWidget(QWidget):
    """
    数据包列表窗口组件，显示捕获的数据包基本信息
    """
    
    # 自定义信号
    packet_selected = pyqtSignal(int)  # 发送选中的数据包索引
    
    def __init__(self, parent=None):
        """
        初始化数据包列表窗口组件
        """
        super().__init__(parent)
        self.packets = []
        self.auto_scroll = True  # 添加自动滚动属性，默认启用
        self._setup_ui()
        
    def _setup_ui(self):
        """
        设置UI组件
        """
        # 创建主布局
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        
        # 创建表格
        self.packet_table = QTableWidget()
        self.packet_table.setColumnCount(6)
        self.packet_table.setHorizontalHeaderLabels(["编号", "时间", "源", "目标", "协议", "长度"])
        
        # 设置表格属性
        self.packet_table.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.packet_table.setSelectionMode(QAbstractItemView.SingleSelection)
        self.packet_table.setSortingEnabled(True)
        self.packet_table.verticalHeader().setVisible(False)
        self.packet_table.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)
        self.packet_table.horizontalHeader().setStretchLastSection(True)
        
        # 设置列宽
        self.packet_table.setColumnWidth(0, 60)  # 编号
        self.packet_table.setColumnWidth(1, 150)  # 时间
        self.packet_table.setColumnWidth(2, 150)  # 源
        self.packet_table.setColumnWidth(3, 150)  # 目标
        self.packet_table.setColumnWidth(4, 80)   # 协议
        
        # 连接信号
        self.packet_table.itemSelectionChanged.connect(self._on_selection_changed)
        self.packet_table.customContextMenuRequested.connect(self._show_context_menu)
        self.packet_table.setContextMenuPolicy(Qt.CustomContextMenu)
        
        # 添加到布局
        layout.addWidget(self.packet_table)
        self.setLayout(layout)
        
    def update_packets(self, packets):
        """
        更新数据包列表
        
        Args:
            packets: 数据包列表
        """
        # 保存滚动位置
        scroll_bar = self.packet_table.verticalScrollBar()
        scroll_position = scroll_bar.value()
        
        # 保存当前选择
        current_row = -1
        selected_items = self.packet_table.selectedItems()
        if selected_items:
            current_row = selected_items[0].row()
            
        # 禁用排序提高性能
        self.packet_table.setSortingEnabled(False)
        
        # 暂停绘制提高性能
        self.packet_table.setUpdatesEnabled(False)
        
        try:
            # 清空表格
            self.packet_table.setRowCount(0)
            
            # 保存新的数据包列表
            self.packets = packets
            
            # 限制最大显示数量以提高性能
            max_display = 10000
            if len(packets) > max_display:
                display_packets = packets[-max_display:]  # 只显示最新的包
            else:
                display_packets = packets
            
            # 设置表格行数
            self.packet_table.setRowCount(len(display_packets))
            
            # 批量添加数据
            for i, packet in enumerate(display_packets):
                self._add_packet_to_table(packet, i)
        except Exception as e:
            print(f"更新数据包表格时出错: {str(e)}")
        finally:
            # 恢复绘制和排序
            self.packet_table.setUpdatesEnabled(True)
            self.packet_table.setSortingEnabled(True)
            
            # 恢复滚动位置
            if scroll_position > 0:
                scroll_bar.setValue(scroll_position)
                
            # 恢复选择
            if current_row >= 0 and current_row < self.packet_table.rowCount():
                self.packet_table.selectRow(current_row)
            
            # 更新状态栏
            if hasattr(self, "statusBar") and self.statusBar:
                self.statusBar.showMessage(f"显示 {self.packet_table.rowCount()} 个数据包")
                
    def _add_packet_to_table(self, packet, row):
        """
        将数据包添加到表格中指定行
        
        Args:
            packet: 数据包对象
            row: 行索引
        """
        try:
            # 编号
            number_item = QTableWidgetItem(str(row + 1))
            number_item.setData(Qt.UserRole, packet)
            self.packet_table.setItem(row, 0, number_item)
            
            # 时间
            time_item = QTableWidgetItem(packet.get("time", ""))
            self.packet_table.setItem(row, 1, time_item)
            
            # 源
            src_item = QTableWidgetItem(packet.get("src", "未知"))
            self.packet_table.setItem(row, 2, src_item)
            
            # 目标
            dst_item = QTableWidgetItem(packet.get("dst", "未知"))
            self.packet_table.setItem(row, 3, dst_item)
            
            # 协议
            protocol_item = QTableWidgetItem(packet.get("protocol", "未知"))
            self.packet_table.setItem(row, 4, protocol_item)
            
            # 长度
            length_item = QTableWidgetItem(str(packet.get("length", "")))
            self.packet_table.setItem(row, 5, length_item)
            
            # 根据协议类型设置行颜色
            self._set_row_color(row, packet.get("protocol", ""))
            
        except Exception as e:
            print(f"将数据包添加到表格时出错 (行 {row}): {str(e)}")
        
    def _set_row_color(self, row, protocol):
        """
        根据协议类型设置行颜色
        
        Args:
            row: 行索引
            protocol: 协议名称
        """
        # 定义协议颜色映射
        protocol_colors = {
            "TCP": QColor(230, 240, 255),      # 浅蓝色
            "UDP": QColor(255, 240, 230),      # 浅橙色
            "ICMP": QColor(255, 230, 230),     # 浅红色
            "HTTP": QColor(230, 255, 230),     # 浅绿色
            "HTTPS": QColor(200, 255, 200),    # 绿色
            "DNS": QColor(255, 255, 230),      # 浅黄色
            "ARP": QColor(240, 230, 255),      # 浅紫色
            "IPv6": QColor(230, 230, 255)      # 浅蓝紫色
        }
        
        # 获取协议颜色，如果没有则使用默认颜色
        color = protocol_colors.get(protocol, None)
        
        if color:
            brush = QBrush(color)
            for col in range(self.packet_table.columnCount()):
                item = self.packet_table.item(row, col)
                if item:
                    item.setBackground(brush)
                    
    def _on_selection_changed(self):
        """
        处理表格选择变化事件
        """
        selected_items = self.packet_table.selectedItems()
        if selected_items:
            # 获取第一列（编号列）的项
            number_item = self.packet_table.item(selected_items[0].row(), 0)
            # 获取存储的原始索引
            packet = number_item.data(Qt.UserRole)
            # 发送信号
            self.packet_selected.emit(self.packets.index(packet))
            
    def _show_context_menu(self, pos):
        """
        显示上下文菜单
        
        Args:
            pos: 鼠标点击位置
        """
        selected_items = self.packet_table.selectedItems()
        if not selected_items:
            return
            
        # 创建上下文菜单
        context_menu = QMenu(self)
        
        # 添加菜单项
        follow_stream_action = QAction("跟踪流", self)
        filter_by_protocol_action = QAction("按协议过滤", self)
        filter_by_source_action = QAction("按源地址过滤", self)
        filter_by_destination_action = QAction("按目标地址过滤", self)
        
        # 添加到菜单
        context_menu.addAction(follow_stream_action)
        context_menu.addSeparator()
        context_menu.addAction(filter_by_protocol_action)
        context_menu.addAction(filter_by_source_action)
        context_menu.addAction(filter_by_destination_action)
        
        # 执行菜单
        context_menu.exec_(self.packet_table.mapToGlobal(pos))
        
    def add_packet(self, packet):
        """
        添加单个数据包到列表末尾
        
        Args:
            packet: 数据包对象
        """
        try:
            # 将数据包添加到列表
            if not hasattr(self, 'packets'):
                self.packets = []
            self.packets.append(packet)
            
            # 限制表格最大行数，防止内存溢出
            max_rows = 10000
            if self.packet_table.rowCount() >= max_rows:
                # 每次达到最大行数时，清除前一半数据
                half_rows = max_rows // 2
                for i in range(half_rows):
                    self.packet_table.removeRow(0)
                
            row = self.packet_table.rowCount()
            self.packet_table.insertRow(row)
            self._add_packet_to_table(packet, row)
            
            # 如果自动滚动启用，滚动到底部
            if self.auto_scroll and row > 2:
                self.packet_table.scrollToItem(self.packet_table.item(row, 0))
                
        except Exception as e:
            print(f"添加数据包到表格时出错: {str(e)}")
            
    def clear(self):
        """
        清空数据包列表
        """
        self.packets = []
        self.packet_table.setRowCount(0) 